import chromadb
from chromadb.config import Settings

class MyVectorDBConnector:
    def __init__(self,collection_name,embedding_fn):
        chroma_client = chromadb.Client(Settings(allow_reset = True))

        chroma_client.reset()

        # 创建一个链接
        self.collection = chroma_client.get_or_create_collection(name="pdf")
        self.embedding_fn = embedding_fn
    
    def add_document(self,documents,metadata={}):
        """向collection中添加文档与向量"""
        self.collection.add(
            embeddings = self.embedding_fn(documents), # 每个文档的向量
            documents = documents, # 文档的原文
            ids = [f"id{i}" for i in range(len(documents))] # 每个文档的 id


        
        )
    def search(self,query,top_n):
        """检索向量数据库"""
        results = self.collection.query(
            query_embeddings=self.embedding_fn([query]),
            n_results=top_n,
        )
        return results



